package com.kamistoat.meimeimall.common.sharding;

import com.kamistoat.meimeimall.common.exception.BizCodeEnum;
import com.kamistoat.meimeimall.common.exception.RRException;
import com.kamistoat.meimeimall.common.sharding.properties.DatabaseProperties;
import com.kamistoat.meimeimall.common.sharding.properties.ShardingProperties;
import com.kamistoat.meimeimall.common.sharding.strategy.AbstractShardingStrategy;
import com.zaxxer.hikari.HikariDataSource;
import org.springframework.boot.autoconfigure.AutoConfigureBefore;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.datasource.LazyConnectionDataSourceProxy;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.annotation.Resource;
import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * 加载分库分表插件
 */
@Configuration
@AutoConfigureBefore(DataSourceAutoConfiguration.class)
@ConditionalOnProperty(name = "sharding.databases", havingValue = "enable")
@EnableConfigurationProperties(ShardingProperties.class)
public class ShardingDataSourceConfiguration {

    @Resource
    private ShardingProperties shardingProperties;

    /**
     * 由于进行SQL重写路由时使用的是创建Statement前阶段的mybatis插件，此时已经获取到了一个真实的数据库连接，无法重新进行分库级别的路由了。
     * 那么只能将路由阶段提前或者将获取数据库连接阶段延后，提前的话就无法拿到待执行的SQL，所以只能选择将获取连接阶段延后。
     * 可以使用spring提供的 {@link LazyConnectionDataSourceProxy}.
     * 使用该延迟类型数据库连接池在创建连接时，并不会真正创建连接，而是产生一个Connection代理类，对该代理Connection进行属性设置，比如开启事务、设置隔离级别等。
     * 配置修改时，只是在内存中进行记录，只有需要执行sql时，才会真正创建一个Connection，并将之前内存中的事务属性、隔离级别属性应用到该真实Connection上。
     * 通过该延迟类型连接池，我们就可以在创建Statement前阶段的mybatis插件中，通过ShardingContext.setShardingDatabase进行分库key的设置，
     * 从而在ShardingDataSource中获取该SQL所属的分库对应的DataSource，来获取对应的Connection。
     */
    @Bean
    public DataSource dataSource() {
        // 创建 LazyConnectionDataSourceProxy
        LazyConnectionDataSourceProxy dataSourceProxy = new LazyConnectionDataSourceProxy();
        // 为这个 LazyConnectionDataSourceProxy 绑定其管理的 DataSource
        dataSourceProxy.setTargetDataSource(shardingDataSource());
        return dataSourceProxy;
    }

    /**
     * 将yaml中的自定义分库分表链接地址解析成为Spring的 DataSource. 并绑定到 LazyConnectionDataSourceProxy 中
     */
    private DataSource shardingDataSource() {
        // 获取yaml中的配置
        Map<String, DatabaseProperties> databasePropertiesMap = shardingProperties.getDatabases();
        if (databasePropertiesMap == null || databasePropertiesMap.size() == 0) {
            throw new RRException(BizCodeEnum.SHARDING_DATABASES_NULL.getMessage(), BizCodeEnum.SHARDING_DATABASES_NULL.getCode());
        }

        Map<Object, Object> targetDataSources = new HashMap<>();
        DataSource dataSource = null;

        /**
         * 解析yaml中的所有总数据源，创建 DataSource
         */
        for (Map.Entry<String, DatabaseProperties> entry : databasePropertiesMap.entrySet()) {
            // 数据源名称
            String dataBaseName = entry.getKey();
            // 数据源属性
            DatabaseProperties databaseProperties = entry.getValue();

            // 获取该数据源的主从库
            Set<Map.Entry<String, Map<Integer, DataSourceProperties>>> entries = databaseProperties.getDataSource().entrySet();
            // 对该数据源主库分库进行处理
            for (Map.Entry<String, Map<Integer, DataSourceProperties>> masterSlave : entries) {
                // 获取主从标识
                String masterSlaveKey = masterSlave.getKey();
                // 获取该数据源主库的所有分库
                Map<Integer, DataSourceProperties> masterSlaveValue = masterSlave.getValue();
                for (Map.Entry<Integer, DataSourceProperties> propertiesEntry : masterSlaveValue.entrySet()) {
                    // 设置分库的标识为 数据源名称+主从标识+分库序号
                    String shardingDataBaseKey = dataBaseName + masterSlaveKey + propertiesEntry.getKey();
                    // 为该分库创建Spring的DataSource
                    dataSource = createDataSource(propertiesEntry.getValue(), HikariDataSource.class);
                    // 添加到Map中
                    targetDataSources.put(shardingDataBaseKey, dataSource);
                }
            }
        }

        if (targetDataSources.size() == 0) {
            throw new RRException(BizCodeEnum.SHARDING_DATABASES_NULL.getMessage(), BizCodeEnum.SHARDING_DATABASES_NULL.getCode());
        }

        if (dataSource == null) {
            throw new RRException(BizCodeEnum.SHARDING_DATABASES_NULL.getMessage(), BizCodeEnum.SHARDING_DATABASES_NULL.getCode());
        }

        /**
         * 将yaml中的配置绑定到分库分表策略抽象类中
         */
        AbstractShardingStrategy.setDatabasePropertiesMap(databasePropertiesMap);

        /**
         * 用于创建LazyConnectionDataSourceProxy时获取真实数据库连接，来获取实际数据库的自动提交配置和隔离级别
         */
        ShardingDataSource shardingDataSource = new ShardingDataSource();
        shardingDataSource.setTargetDataSources(targetDataSources);
        shardingDataSource.setDefaultTargetDataSource(dataSource);
        shardingDataSource.setLenientFallback(false);
        shardingDataSource.afterPropertiesSet();
        return shardingDataSource;
    }

    /**
     * 根据Spring的数据源链接配置，创建Spring的DataSource
     *
     * @param properties 链接配置
     * @param type       Spring DataSource种类，默认是 HikariDataSource.class
     * @param <T>
     * @return
     */
    @SuppressWarnings("unchecked")
    private <T> T createDataSource(DataSourceProperties properties, Class<? extends DataSource> type) {
        return (T) properties.initializeDataSourceBuilder().type(type).build();
    }

    /**
     * 继承 Spring内部类 AbstractRoutingDataSource。
     * 数据源路由器类。
     * 所有的数据源链接参数保存在自身 .targetDataSource 这个Map属性中。
     * 每次建立数据源链接时，根据determineCurrentLookupKey() 方法的返回值，去 .targetDataSource 寻找对应的链接参数建立链接。
     */
    public static class ShardingDataSource extends AbstractRoutingDataSource {
        /**
         * 确定数据源标识，即将建立链接。直接从 ShardingContext 中把拦截器确定的标识拿出来即可。
         */
        @Override
        protected Object determineCurrentLookupKey() {
            return ShardingContext.getShardingDatabase();
        }
    }
}
